// -----------------------------------------------------------------------------
// Nx1 A R I T I Nx1 E   R E S E A R C H   I Nx2 S T I T U T E   Nx2 E T H E R L A Nx2 D S
// -----------------------------------------------------------------------------
// Copyright (C) April 27, 2010 MARIN - All rights reserved - http://www.marin.nl
// -----------------------------------------------------------------------------
// Program    : mermaid
// Module     : vbm
// File     : IterCudaSolver.cpp
// -----------------------------------------------------------------------------

#include <iterative-cuda.hpp>
#include "defines.h"
#include "IterCudaSolver.h"
#include "assertplus.h"

#define MATRIXMARKETFILE "matrix.txt"
#define TEMP             "temp.txt"

// -----------------------------------------------------------------------------
// function IterCudaSolver::IterCudaSolver
// -----------------------------------------------------------------------------
/// constructor
/** detailed description */
// -----------------------------------------------------------------------------
IterCudaSolver::IterCudaSolver(Array2D<REAL> *pcc,
                               Array2D<REAL> *pcs,
                               Array2D<REAL> *pcw,
                               int Nx1,
                               int Nx2,
                               REAL tolerance)
: base(Nx1, Nx2, tolerance)
, m_pCpuMatrix(0)
, m_pGpuMatrix(0)
, m_pInvDiagGpu(0)
, m_pPermInvDiagGpu(0)
, m_pDiagPre(0)
, m_pB(0)
, m_pX(0)
, m_n(0)
{
    preconditioner(*pcc, *pcs, *pcw);
}


// -----------------------------------------------------------------------------
// function IterCudaSolver::~IterCudaSolver()
// -----------------------------------------------------------------------------
/// destructor
/** detailed description */
// -----------------------------------------------------------------------------
IterCudaSolver::~IterCudaSolver()
{
#ifdef DEBUG
    delete m_pCpuMatrix;
#endif
    delete m_pGpuMatrix;
    delete m_pInvDiagGpu;
    delete m_pPermInvDiagGpu;
    delete m_pDiagPre;
    delete m_pB;
    delete m_pX;
}

void IterCudaSolver::preconditioner(Array2D<REAL>& matc,
                                    Array2D<REAL>& mats,
                                    Array2D<REAL>& matw)
{
    // TODO add code to iterative cuda library to avoid having to write a temporary matrix market file
    makeMatrixMarketFile(matc, mats, matw, MATRIXMARKETFILE);

    // create sparse matrix
    m_pCpuMatrix = CpuMatrixType::read_matrix_market_file(MATRIXMARKETFILE);

    // diagonal preconditioner
    m_n = m_pCpuMatrix->row_count();
    REAL *diag = new REAL[m_n];
    REAL *invDiag = new REAL[m_n];
    m_pCpuMatrix->extract_diagonal(diag);
    for (unsigned int i = 0; i < m_n; ++i)
    {
        invDiag[i] = 1./diag[i];
    }
    m_pInvDiagGpu = new GpuVectorType(invDiag, m_n);
    delete[] invDiag;
    delete[] diag;

    // copy matrix to GPU
    m_pGpuMatrix = new GpuMatrixType(*m_pCpuMatrix);
    m_pPermInvDiagGpu = new GpuVectorType(m_n);
    m_pGpuMatrix->permute(*m_pPermInvDiagGpu, *m_pInvDiagGpu);
    m_pDiagPre = new iterative_cuda::diagonal_preconditioner<GpuVectorType>(*m_pPermInvDiagGpu);

    // allocate vectors
    m_pB = new REAL[m_n];
    m_pX = new REAL[m_n];

    // adjust tolerance
    if (sizeof(REAL) < sizeof(double))
    {
        m_tolerance = 1e-4;
    }
    else
    {
        m_tolerance = 1e-12;
    }

#ifndef DEBUG
    delete m_pCpuMatrix;
#endif
}


void IterCudaSolver::solve(Array2D<REAL> *pB,
                           Array2D<REAL> *pX)
{
    // syntactic sugar
    Array2D<REAL>& bpsi = *pB;
    Array2D<REAL>& psi  = *pX;

    // this is done inside the WavesComputer::compute() parallelization
#ifdef openmp1
#pragma omp barrier
#endif
    {
        // convert bpsi to vector b
        int k = 0;
        for (int j = 1; j <= m_Nx2; ++j)
        {
            for (int i = 1; i <= m_Nx1; ++i)
            {
                *(m_pB + k) = bpsi[i][j];
                ++k;
            }
        }

        GpuVectorType bGpu(m_pB, m_n);
        GpuVectorType bPermGpu(m_n);
        m_pGpuMatrix->permute(bPermGpu, bGpu);
        GpuVectorType xPermGpu(m_n);
        xPermGpu.fill(0);

        // perform solve
        unsigned it_count;
        iterative_cuda::gpu_cg(*m_pGpuMatrix, *m_pDiagPre, xPermGpu, bPermGpu, m_tolerance, /*max_iterations*/ 0, &it_count);
        GpuVectorType xGpu(m_n);
        m_pGpuMatrix->unpermute(xGpu, xPermGpu);
        xGpu.to_cpu(m_pX);

        // convert result vector x to psi
        k = 0;
        for (int j = 1; j <= m_Nx2; ++j)
        {
            for (int i = 1; i <= m_Nx1; ++i)
            {
                psi[i][j] = *(m_pX + k);
                ++k;
            }
        }
    } // master openmp1
#ifdef openmp1
#pragma omp barrier
#endif
}


void IterCudaSolver::makeMatrixMarketFile(Array2D<REAL>& matc,
                                          Array2D<REAL>& mats,
                                          Array2D<REAL>& matw,
                                          const char* file)
{
    FILE* fp = fopen(TEMP, "w");
    int count = 0;
    for (int i = 1; i <= m_Nx1; ++i) // east - west direction
    {
        for (int j = 1; j <= m_Nx2; ++j) // north - south direction
        {
            int k = m_Nx1 * (j - 1) + (i - 1);
            fprintf(fp, "%d %d %G\n", k + 1, k + 1, matc[i][j]);
            ++count;
            if (k - 1 >= 0)
            {
                fprintf(fp, "%d %d %G\n", k - 1 + 1, k + 1, matw[i][j]);
                fprintf(fp, "%d %d %G\n", k + 1, k - 1 + 1, matw[i][j]);
                ++count;
            }
            if (k - m_Nx1 >= 0)
            {
                fprintf(fp, "%d %d %G\n", k - m_Nx1 + 1, k + 1, mats[i][j]);
                fprintf(fp, "%d %d %G\n", k + 1, k - m_Nx1 + 1, mats[i][j]);
                ++count;
            }
        }
    }
    fclose(fp);
    fp = fopen(file, "w");
    fprintf(fp, "%%%%MatrixMarket matrix coordinate real general\n");
    fprintf(fp, "%d %d %d\n", m_Nx1 * m_Nx2, m_Nx1 * m_Nx2,
            m_Nx1 * m_Nx2 + 2 * (m_Nx1 * m_Nx2 - 1) + 2 * (m_Nx1 * m_Nx2 - m_Nx1));
    fclose(fp);
    char command[500];
    sprintf(command, "sort -k 1n -k 2n %s >> %s", TEMP, file);
    int ignore = 0;
    ignore = system(command);
    unlink(TEMP);
}
